import math

import pandas as pd
import numpy as np
import copy


def EMA(s: pd.Series, N):
    """
    EMA(df, n,priceSign='close'):
    #Exponential Moving Average
    EMA是指数平滑移动平均线，也叫EXPMA指标，也称为：SMMA
    是平均线的一个变种，EMA均线较MA更加专业一些。

    """
    return pd.Series.ewm(s, span=N, min_periods=1, adjust=True).mean()


def MA(s: pd.Series, N):
    return pd.Series.rolling(s, N, min_periods=1).mean()


def SMA(s: pd.Series, N, M=1):
    """
    威廉SMA算法

    本次修正主要是对于返回值的优化,现在的返回值会带上原先输入的索引index
    威廉SMA  参考https://www.joinquant.com/post/867
    """
    ret = []
    i = 0
    length = len(s)
    # 跳过X中前面几个 nan 值
    while i < length:
        if np.isnan(s.iloc[i]):
            i += 1
        else:
            break
    preY = s.iloc[i]  # Y'
    ret.append(preY)
    while i < length:
        Y = (M * s.iloc[i] + (N - M) * preY) / float(N)
        ret.append(Y)
        preY = Y
        i += 1
    return pd.Series(ret, index=s.tail(len(ret)).index)


def AM(s: pd.Series, af=MA, N=0):
    """
    自适应均值算法：Adaptive Mean
    :params af(average function): 是均值的计算方法，MA，SMA，EMA
    :params N:  当N=0时，用整个序列的长度计算均值

    1. 计算af:
       a=af(s,N)

    """
    dl = len(s)
    if N == 0:
        N = dl
    a = af(s, N)
    i = 0
    j = 0
    am = copy.deepcopy(a)
    at = []
    ud = []
    idx = 0
    for i in range(1, dl):
        uds = SINGLE_CROSS(s.iloc[:i + 1], a.iloc[:i + 1])
        if uds == 1:
            idx = s.iloc[j:i].idxmin()
            ud.append([uds, idx, s.index[i], s[idx], s.iloc[i]])
            j = len(s[:idx]) - 1
        elif uds == -1:
            idx = s.iloc[j:i].idxmax()
            ud.append([uds, idx, s.index[i], s[idx], s.iloc[i]])
            j = len(s[:idx]) - 1
        if j != 0:
            at = af(s.iloc[j:i], N)


def DIFF(s: pd.Series, N=1):
    return pd.Series(s).diff(N)


def HHV(s: pd.Series, periods, center=0):
    """
    取得 data 在 periods 周期内的最大值，当 periods 为偶数且 center==0 时，前边取一半（不含当前位置），
    :param Series:pd.Series
    :param periods:
    :param center: 0==center, 1==before, -1=after
    :return:
    """
    if center == 0:
        hh = s.rolling(window=periods, min_periods=1, center=True).max()
    elif center == 1:
        hh = s.rolling(window=periods, min_periods=1, center=False).max()
    else:  # center==-1:
        hh = s[::-1].rolling(window=periods, min_periods=1, center=False).max()[::-1]
    return hh


def LLV(s: pd.Series, periods, center=0):
    """
    取得 data 在 periods 周期内的最小值，当 periods 为偶数且 center==0 时，前边取一半（不含当前位置），
    :param Series:pd.Series
    :param periods:
    :param center: 0==center, 1==before, -1=after
    :return:
    """
    if center == 0:
        ll = s.rolling(window=periods, min_periods=1, center=True).min()
    elif center == 1:
        ll = s.rolling(window=periods, min_periods=1, center=False).min()
    else:  # center==-1:
        ll = s[::-1].rolling(window=periods, min_periods=1, center=False).min()[::-1]
    return ll


def SUM(s: pd.Series, N):
    return pd.Series.rolling(s, N).sum()


def ABS(s: pd.Series):
    return abs(s)


def ATR(s: pd.Series, n, m=1):
    '''
    def ATR(df, n,ksgn='close'):
    #ATR,均幅指标（Average True Ranger）,取一定时间周期内的股价波动幅度的移动平均值，主要用于研判买卖时机

    【输入】
        s, pd.Series 格式数据源
        n，时间长度
    【输出】
        df, pd.Series 格式数据源,
        增加了一栏：atr_{n}，输出数据
     该指标价值越高，趋势改变的可能性就越高；该指标的价值越低，趋势的移动性就越弱。
    [14,14,100,100]
    '''
    xnam = 'ATR'

    hh = s.rolling(window=m, min_periods=1).max()
    ll = s.rolling(window=m, min_periods=1).min()
    tr_srs = hh - ll
    df = pd.DataFrame(tr_srs, index=s.index, columns=['TR'])
    ds = tr_srs.ewm(span=n + m, min_periods=1, adjust=True, ignore_na=False).mean()
    ds.index = df.index
    df[xnam] = ds
    # print('ds',ds.head())
    return df


def TR(df, sgn_close='close', sgn_high='high', sgn_low='low', n=5):
    """
    计算真实波幅
    【输入】
        df, pd.DataFrame 格式数据源
        sgn_close，列名，一般是：close收盘价
    【输出】
        pd.DataFrame格式数据源,

    """

    i = 0
    TR_l = [df[sgn_high][i] - df[sgn_low][0]]
    while i < len(df) - 1:  # df.index[-1]:
        TR = max(df[sgn_high].iloc[i + 1], df[sgn_close].iloc[i]) - min(df[sgn_low].iloc[i + 1], df[sgn_close].iloc[i])
        TR_l.append(TR)
        i += 1
    TR_s = pd.Series(TR_l, index=df.index, name='TR')
    ATR = pd.Series(TR_s.ewm(span=n, min_periods=n).mean(), name='ATR', index=df.index)
    return TR_s, ATR


def MAX(A: pd.Series, B: pd.Series):
    """

    """
    var = IF(A > B, A, B)
    return var


def MIN(A: pd.Series, B: pd.Series):
    var = IF(A < B, A, B)
    return var


def RSI(s: pd.Series, n):
    '''
    def rsi(df, n):
    相对强弱指标,Relative Strength Index、相对力度指数
    RSI，是通过比较一段时期内的平均收盘涨数和平均收盘跌数来分析市场买沽盘的意向和实力，从而作出未来市场的走势。
    RSI通过特定时期内股价的变动情况计算市场买卖力量对比，来判断股票价格内部本质强弱、推测价格未来的变动方向的技术指标。

    【输入】
        s, pd.Series 格式数据源
        n，时间长度,一般为14
    【输出】
        srs, pd.Series 格式数据源,
        增加了一栏：rsi，输出数据
    '''
    xnam = 'RSI'
    i = 0
    UpI = [0]
    DoI = [0]
    while i + 1 <= len(s) - 1:  # df.index[-1]
        UpMove = s.iloc[i + 1] - s.iloc[i]
        DoMove = s.iloc[i] - s.iloc[i + 1]

        # Range=abs(df['high'].iloc[i+1]-df['low'].iloc[i])-abs(df['low'].iloc[i+1]-df['high'].iloc[i])
        if UpMove > DoMove and UpMove > 0:
            UpD = UpMove
        else:
            UpD = 0
        UpI.append(UpD)
        if DoMove > UpMove and DoMove > 0:
            DoD = DoMove
        else:
            DoD = 0
        DoI.append(DoD)
        i = i + 1
    UpI = pd.Series(UpI)
    DoI = pd.Series(DoI)
    PosDI = UpI.ewm(span=n, min_periods=n - 1, adjust=True, ignore_na=False).mean()
    NegDI = DoI.ewm(span=n, min_periods=n - 1, adjust=True, ignore_na=False).mean()

    ds = pd.Series(PosDI / (PosDI + NegDI), name=xnam)
    ds.index = s.index
    srs = ds * 100.
    return srs


def SINGLE_CROSS(A: pd.Series, B: pd.Series):
    """
    A、B：pd.Series
    Return:
    如果A[-1]上穿B[-1]：1
    如果A[-1]下穿B[-1]：-1
    其它：0
    """

    if A.iloc[-2] < B.iloc[-2] and A.iloc[-1] > B.iloc[-1]:
        return 1
    elif A.iloc[-2] > B.iloc[-2] and A.iloc[-1] < B.iloc[-1]:
        return -1
    else:
        return 0


def CROSS(A: pd.Series, B: pd.Series):
    """A<B then A>B  A上穿B B下穿A

    Arguments:
        A: pd.Series
        B: Pd.Series

    Returns:
       pd.s, where 1==a up_cross b, -1 == a down_cross b
    """

    var = np.where(A > B, 1, 0)
    crss = pd.Series(var, index=A.index).diff(1)
    crss[0] = 0
    return crss  # .apply(int)


def COUNT(COND, N):
    """
    2018/05/23 修改

    参考https://github.com/QUANTAXIS/QUANTAXIS/issues/429

    现在返回的是series
    """
    return pd.Series(np.where(COND, 1, 0), index=COND.index).rolling(N).sum()


def IF(COND, V1, V2):
    """如果条件 COND 成立，返回 V1 ， 否则，返回 V2"""
    var = np.where(COND, V1, V2)
    return pd.Series(var, index=V1.index)


def IFAND(COND1, COND2, V1, V2):
    """如果条件 COND1 和 COND2 同时成立，返回 V1 ， 否则，返回 V2"""
    var = np.where(np.logical_and(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def IFOR(COND1, COND2, V1, V2):
    """如果条件 COND1 和 COND2 中有一个成立，返回 V1 ， 否则，返回 V2"""
    var = np.where(np.logical_or(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def REF(s: pd.Series, N):
    var = s.diff(N)
    var = s - var
    return var


def LAST(COND, N1, N2):
    """表达持续性
    从前N1日到前N2日一直满足COND条件

    Arguments:
        COND {[type]} -- [description]
        N1 {[type]} -- [description]
        N2 {[type]} -- [description]
    """
    N2 = 1 if N2 == 0 else N2
    assert N2 > 0
    assert N1 > N2
    return COND.iloc[-N1:-N2].all()


def STD(s: pd.Series, N):
    return pd.Series.rolling(s, N).std()


def MADEV(s: pd.Series, N):
    """
    平均绝对偏差 mean absolute deviation
    abs(x-x.mean()).mean()
    """
    return s.rolling(N).apply(lambda x: (np.abs(x - x.mean())).mean(), raw=True)


def AVEDEV(Series, N):
    """
    平均绝对偏差 mean absolute deviation
    修正: 2018-05-25

    之前用mad的计算模式依然返回的是单值
    """
    return Series.rolling(N).apply(lambda x: (np.abs(x - x.mean())).mean(), raw=True)


def MACD(s: pd.Series, FAST, SLOW, MID):
    """macd指标 仅适用于Series
    对于DATAFRAME的应用请使用QA_indicator_macd
    """
    EMAFAST = EMA(s, FAST)
    EMASLOW = EMA(s, SLOW)
    DIFF = EMAFAST - EMASLOW
    DEA = EMA(DIFF, MID)
    MACD = (DIFF - DEA) * 2
    DICT = {'DIFF': DIFF, 'DEA': DEA, 'MACD': MACD}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBIBOLL(s, N1, N2, N3, N4, N, M):  # 多空布林线

    bbiboll = BBI(s, N1, N2, N3, N4)
    UPER = bbiboll + M * STD(bbiboll, N)
    DOWN = bbiboll - M * STD(bbiboll, N)
    DICT = {'BBIBOLL': bbiboll, 'UPER': UPER, 'DOWN': DOWN}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBI(s, N1, N2, N3, N4):
    '多空指标'

    bbi = (MA(s, N1) + MA(s, N2) +
           MA(s, N3) + MA(s, N4)) / 4
    DICT = {'BBI': bbi}
    VAR = pd.DataFrame(DICT)
    return VAR


def ACSP(s: pd.s, spp=5, acp=5):
    """速度加速度
    速度：价格的离差
    加速度：速度的离差
    :return :速度、加速度序列
    """
    sp = s.diff(spp)
    ac = sp.diff(acp)
    return sp, ac


def BARLAST(cond, yes=True):
    """支持MultiIndex的cond和DateTimeIndex的cond
    条件成立  yes= True 或者 yes=1 根据不同的指标自己定

    Arguments:
        cond {[type]} -- [description]
    """
    if isinstance(cond.index, pd.MultiIndex):
        return len(cond) - cond.index.levels[0].tolist().index(cond[cond != yes].index[-1][0]) - 1
    elif isinstance(cond.index, pd.DatetimeIndex):
        return len(cond) - cond.index.tolist().index(cond[cond != yes].index[-1]) - 1


ma = MA
ema = EMA
sma = SMA
diff = DIFF
hhv = HHV
llv = LLV
sum = SUM
cross = CROSS
single_cross = SINGLE_CROSS
count = COUNT
ifand = IFAND
ifor = IFOR
ref = REF
last = LAST
std = STD
avedev = AVEDEV
macd = MACD
bbiboll = BBIBOLL
bbi = BBI
barlast = BARLAST

XARROUND = lambda x, y: np.round(y * (round(x / y - math.floor(x / y) + 0.00000000001) + math.floor(x / y)), 2)


def SGN(s: pd.Series):
    sg = s.apply(lambda x: 1 if x > 0 else -1 if x < 0 else 0)
    return sg


def draw_back(s: pd.s, is_cusum=True):
    """回撤计算：
    :param s: 收益率（收益）序列
    :param is_cusum: 是否是累积收益（率）
    :return : 
    """
    dl = len(s)
    if is_cusum:
        hs = s.rolling(window=dl, min_periods=1).max()
        ls = s.rolling(window=dl, min_periods=1).min()
        r = hs - ls
    else:
        hs = s.cumsum().rolling(window=dl, min_periods=1).max()
        ls = s.cumsum().rolling(window=dl, min_periods=1).min()
        r = hs - ls
    return r
